/* -*- Mode: java; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 4 -*-
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package org.mozilla.javascript.optimizer;

import static org.mozilla.classfile.ClassFileWriter.ACC_PUBLIC;
import static org.mozilla.classfile.ClassFileWriter.ACC_STATIC;

import java.lang.invoke.MethodType;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.mozilla.classfile.ByteCode;
import org.mozilla.classfile.ClassFileWriter;
import org.mozilla.javascript.CompilerEnvirons;
import org.mozilla.javascript.IRFactory;
import org.mozilla.javascript.JSCode;
import org.mozilla.javascript.JSDescriptor;
import org.mozilla.javascript.JavaAdapter;
import org.mozilla.javascript.Parser;
import org.mozilla.javascript.ScriptRuntime;
import org.mozilla.javascript.ast.AstRoot;
import org.mozilla.javascript.ast.FunctionNode;
import org.mozilla.javascript.ast.ScriptNode;

/**
 * Generates class files from script sources.
 *
 * <p>since 1.5 Release 5
 *
 * @author Igor Bukanov
 */
public class ClassCompiler {
    /**
     * Construct ClassCompiler that uses the specified compiler environment when generating classes.
     */
    public ClassCompiler(CompilerEnvirons compilerEnv) {
        if (compilerEnv == null) throw new IllegalArgumentException();
        this.compilerEnv = compilerEnv;
        this.mainMethodClassName = Codegen.DEFAULT_MAIN_METHOD_CLASS;
    }

    /**
     * Set the class name to use for main method implementation. The class must have a method
     * matching {@code public static void main(Script sc, String[] args)}, it will be called when
     * {@code main(String[] args)} is called in the generated class. The class name should be fully
     * qualified name and include the package name like in {@code org.foo.Bar}.
     */
    public void setMainMethodClass(String className) {
        // XXX Should this check for a valid class name?
        mainMethodClassName = className;
    }

    /**
     * Get the name of the class for main method implementation.
     *
     * @see #setMainMethodClass(String)
     */
    public String getMainMethodClass() {
        return mainMethodClassName;
    }

    /** Get the compiler environment the compiler uses. */
    public CompilerEnvirons getCompilerEnv() {
        return compilerEnv;
    }

    /** Get the class that the generated target will extend. */
    public Class<?> getTargetExtends() {
        return targetExtends;
    }

    /**
     * Set the class that the generated target will extend.
     *
     * @param extendsClass the class it extends
     */
    public void setTargetExtends(Class<?> extendsClass) {
        targetExtends = extendsClass;
    }

    /** Get the interfaces that the generated target will implement. */
    public Class<?>[] getTargetImplements() {
        return targetImplements == null ? null : targetImplements.clone();
    }

    /**
     * Set the interfaces that the generated target will implement.
     *
     * @param implementsClasses an array of Class objects, one for each interface the target will
     *     extend
     */
    public void setTargetImplements(Class<?>[] implementsClasses) {
        targetImplements = implementsClasses == null ? null : implementsClasses.clone();
    }

    /**
     * Build class name for an auxiliary class generated by compiler. If the compiler needs to
     * generate extra classes beyond the main class, it will call this function to build the
     * auxiliary class name. The default implementation simply appends auxMarker to mainClassName
     * but this can be overridden.
     */
    protected String makeAuxiliaryClassName(String mainClassName, String auxMarker) {
        return mainClassName + auxMarker;
    }

    /**
     * Compile JavaScript source into one or more Java class files. The first compiled class will
     * have name mainClassName. If the results of {@link #getTargetExtends()} or {@link
     * #getTargetImplements()} are not null, then the first compiled class will extend the specified
     * super class and implement specified interfaces.
     *
     * @return array where elements with even indexes specifies class name and the following odd
     *     index gives class file body as byte[] array. The initial element of the array always
     *     holds mainClassName and array[1] holds its byte code.
     */
    public Object[] compileToClassFiles(
            String source, String sourceLocation, int lineno, String mainClassName) {
        Parser p = new Parser(compilerEnv);
        AstRoot ast = p.parse(source, sourceLocation, lineno);
        IRFactory irf = new IRFactory(compilerEnv, source);
        ScriptNode tree = irf.transformTree(ast);

        if (compilerEnv.isGeneratingSource()) {
            tree.setRawSource(source);
            tree.setRawSourceBounds(0, source.length());
        }

        // release reference to original parse tree & parser
        irf = null;
        ast = null;
        p = null;

        Class<?> superClass = getTargetExtends();
        Class<?>[] interfaces = getTargetImplements();
        String scriptClassName;
        boolean isPrimary = (interfaces == null && superClass == null);
        if (isPrimary) {
            scriptClassName = mainClassName;
        } else {
            scriptClassName = makeAuxiliaryClassName(mainClassName, "1");
        }

        Codegen codegen = new Codegen();
        codegen.setMainMethodClass(mainMethodClassName);
        JSDescriptor.Builder builder = new JSDescriptor.Builder();
        OptJSCode.BuilderEnv builderEnv = new OptJSCode.BuilderEnv(scriptClassName);
        byte[] scriptClassBytes =
                codegen.compileToClassFile(
                        compilerEnv, builder, builderEnv, scriptClassName, tree, source, false);
        Object[] auxilaryClasses = buildDescriptorsAndMain(scriptClassName, builder);
        if (isPrimary) {
            var result = new Object[auxilaryClasses.length + 2];
            System.arraycopy(auxilaryClasses, 0, result, 2, auxilaryClasses.length);
            result[0] = scriptClassName;
            result[1] = scriptClassBytes;
            return result;
        }
        int functionCount = tree.getFunctionCount();
        HashMap<String, Integer> functionNames = new HashMap<>();
        for (int i = 0; i != functionCount; ++i) {
            FunctionNode ofn = tree.getFunctionNode(i);
            String name = ofn.getName();
            if (name != null && name.length() != 0) {
                functionNames.put(name, ofn.getParamCount());
            }
        }
        if (superClass == null) {
            superClass = ScriptRuntime.ObjectClass;
        }
        byte[] mainClassBytes =
                JavaAdapter.createAdapterCode(
                        functionNames, mainClassName, superClass, interfaces, scriptClassName);

        var result = new Object[auxilaryClasses.length + 4];
        System.arraycopy(auxilaryClasses, 0, result, 4, auxilaryClasses.length);
        result[0] = mainClassName;
        result[1] = mainClassBytes;
        result[2] = scriptClassName;
        result[3] = scriptClassBytes;
        return result;
    }

    /**
     * Build script class initialisation method and main method. The init method will create all the
     * class descriptors, and the main method will create a {@link JSDescriptor} object based on the
     * first descriptor and pass that to the main method in the runtime.
     */
    private Object[] buildDescriptorsAndMain(String mainClassName, JSDescriptor.Builder builder) {
        var classes = new HashMap<String, byte[]>();
        var mainName = mainClassName + "Main";

        var cfw = new ClassFileWriter(mainName, "java.lang.Object", "");
        var builders = new ArrayList<JSDescriptor.Builder<?>>();
        buildDescriptor(cfw, builder, classes, builders);
        cfw.startMethod("<clinit>", "()V", ACC_STATIC);
        cfw.addLoadConstant(builders.size());
        cfw.add(ByteCode.ANEWARRAY, "org/mozilla/javascript/JSDescriptor");
        for (var b : builders) {
            int index = ((OptJSCode.Builder) b.code).index;
            int parent = b.parent == null ? 0 : ((OptJSCode.Builder) b.parent.code).index;
            populateDescriptorEntry(cfw, index, parent);
        }
        cfw.add(
                ByteCode.PUTSTATIC,
                mainClassName,
                Codegen.DESCRIPTORS_FIELD_NAME,
                Codegen.DESCRIPTORS_FIELD_SIGNATURE);
        cfw.add(ByteCode.RETURN);
        cfw.stopMethod(0);

        cfw.startMethod("main", "([Ljava/lang/String;)V", (short) (ACC_STATIC | ACC_PUBLIC));
        cfw.add(ByteCode.NEW, "org.mozilla.javascript.JSScript");
        cfw.add(ByteCode.DUP);
        cfw.add(
                ByteCode.GETSTATIC,
                mainClassName,
                Codegen.DESCRIPTORS_FIELD_NAME,
                Codegen.DESCRIPTORS_FIELD_SIGNATURE);
        cfw.addLoadConstant(0);
        cfw.add(ByteCode.AALOAD);
        cfw.add(ByteCode.ACONST_NULL);
        cfw.addInvoke(
                ByteCode.INVOKESPECIAL,
                "org/mozilla/javascript/JSScript",
                "<init>",
                "(Lorg/mozilla/javascript/JSDescriptor;Lorg/mozilla/javascript/Scriptable;)V");
        cfw.add(ByteCode.ALOAD_0);
        // Call mainMethodClass.main(Script script, String[] args)
        cfw.addInvoke(
                ByteCode.INVOKESTATIC,
                mainMethodClassName,
                "main",
                "(Lorg/mozilla/javascript/Script;[Ljava/lang/String;)V");
        cfw.add(ByteCode.RETURN);
        cfw.stopMethod(1);

        var result = new Object[classes.size() * 2 + 2];
        int count = 0;
        result[count++] = mainName;
        result[count++] = cfw.toByteArray();
        for (var e : classes.entrySet()) {
            result[count++] = e.getKey();
            result[count++] = e.getValue();
        }
        return result;
    }

    /**
     * Generates a byte code for a method to create a {@link JSDescriptor}. It might be more
     * efficient in future to use a const dynamic for this, when our bytecode generation and
     * supported platforms can all handle it.
     */
    private void buildDescriptor(
            ClassFileWriter cfw,
            JSDescriptor.Builder builder,
            Map<String, byte[]> classes,
            List<JSDescriptor.Builder<?>> builders) {
        builders.add(builder);
        cfw.startMethod(
                "init" + functionId(builder),
                "(Lorg/mozilla/javascript/JSDescriptor;)Lorg/mozilla/javascript/JSDescriptor;",
                ACC_STATIC);

        cfw.add(ByteCode.NEW, "org.mozilla.javascript.JSDescriptor");
        cfw.add(ByteCode.DUP);
        buildCode(cfw, builder.code, classes);
        buildCode(cfw, builder.constructor, classes);
        cfw.addALoad(0);
        cfw.addLoadConstant(builder.paramAndVarNames.length);
        cfw.add(ByteCode.ANEWARRAY, "java/lang/String");
        if (builder.paramAndVarNames.length > 0) {
            cfw.addLoadConstant(0);
            cfw.addIStore(2);
            for (int p = 0; p < builder.paramAndVarNames.length; p++) {
                cfw.add(ByteCode.DUP);
                cfw.addILoad(2);
                cfw.addLoadConstant(builder.paramAndVarNames[p]);
                cfw.add(ByteCode.AASTORE);
                cfw.add(ByteCode.IINC, 2, 1);
            }
        }
        cfw.addLoadConstant(builder.paramIsConst.length);
        cfw.add(ByteCode.NEWARRAY, ByteCode.T_BOOLEAN);
        if (builder.paramIsConst.length > 0) {
            cfw.addLoadConstant(0);
            cfw.addIStore(2);
            for (int p = 0; p < builder.paramAndVarNames.length; p++) {
                cfw.add(ByteCode.DUP);
                cfw.addILoad(2);
                cfw.add(builder.paramIsConst[p] ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
                cfw.add(ByteCode.BASTORE);
                cfw.add(ByteCode.IINC, 2, 1);
            }
        }
        cfw.add(builder.isStrict ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.isScript ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.isTopLevel ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.isES6Generator ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.isShorthand ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.hasPrototype ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.hasLexicalThis ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.isEvalFunction ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.hasRestArg ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.addLoadConstant(builder.sourceFile);
        cfw.addLoadConstant(builder.rawSource);
        cfw.addLoadConstant(builder.rawSourceStart);
        cfw.addLoadConstant(builder.rawSourceEnd);
        if (builder.name == null) {
            cfw.add(ByteCode.ACONST_NULL);
        } else {
            cfw.addLoadConstant(builder.name);
        }
        cfw.addLoadConstant(builder.languageVersion);
        cfw.addLoadConstant(builder.paramAndVarCount);
        cfw.addLoadConstant(builder.paramCount);
        cfw.addLoadConstant(builder.arity);
        cfw.add(builder.hasDefaultParameters ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.requiresActivationFrame ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.requiresArgumentObject ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(builder.declaredAsFunctionExpression ? ByteCode.ICONST_1 : ByteCode.ICONST_0);
        cfw.add(ByteCode.ACONST_NULL);
        cfw.add(ByteCode.ACONST_NULL);
        cfw.addLoadConstant(builder.functionType);
        // Call the constructor
        var type =
                MethodType.methodType(
                                void.class,
                                JSDescriptor.class.getConstructors()[0].getParameterTypes())
                        .toMethodDescriptorString();
        cfw.addInvoke(
                ByteCode.INVOKESPECIAL, "org.mozilla.javascript.JSDescriptor", "<init>", type);
        cfw.addAStore(1);
        cfw.addALoad(1);
        cfw.add(ByteCode.DUP);
        cfw.addLoadConstant(builder.nestedFunctions.size());
        cfw.add(ByteCode.ANEWARRAY, "org/mozilla/javascript/JSDescriptor");
        if (builder.nestedFunctions.size() > 0) {
            cfw.addLoadConstant(0);
            cfw.addIStore(2);
            // Loop over the nested functions adding them to the list
            for (var child : builder.nestedFunctions) {
                cfw.add(ByteCode.DUP);
                cfw.addILoad(2);
                cfw.addALoad(1);
                cfw.addInvoke(
                        ByteCode.INVOKESTATIC,
                        cfw.getClassName(),
                        "init" + functionId((JSDescriptor.Builder) child),
                        "(Lorg/mozilla/javascript/JSDescriptor;)Lorg/mozilla/javascript/JSDescriptor;");
                cfw.add(ByteCode.AASTORE);
                cfw.add(ByteCode.IINC, 2, 1);
            }
        }
        cfw.addInvoke(
                ByteCode.INVOKESTATIC,
                "org.mozilla.javascript.optimizer.OptRuntime",
                "listOf",
                "([Ljava/lang/Object;)Ljava/util/List;");
        cfw.add(
                ByteCode.PUTFIELD,
                "org.mozilla.javascript.JSDescriptor",
                "nestedFunctions",
                "Ljava/util/List;");

        cfw.add(ByteCode.ARETURN);
        cfw.stopMethod(2);
        for (var child : builder.nestedFunctions) {
            buildDescriptor(cfw, (JSDescriptor.Builder) child, classes, builders);
        }
    }

    /**
     * Generates a subclass of {@link OptJSCode} to invoke a compiled method, and the code to
     * instantiate that class.
     */
    private void buildCode(
            ClassFileWriter cfw, JSCode.Builder builder, Map<String, byte[]> classes) {
        if (builder instanceof JSCode.NullBuilder<?>) {
            cfw.add(ByteCode.ACONST_NULL);
        } else {
            var code = (OptJSCode.Builder<?>) builder;
            code.buildByteCode(cfw);
            classes.put(code.getClassName(), code.getClassBytes());
        }
    }

    private void populateDescriptorEntry(ClassFileWriter cfw, int index, int parent) {
        // Start with the array on top.
        cfw.add(ByteCode.DUP); // array, array.
        cfw.add(ByteCode.DUP); // array, array, array.
        cfw.addLoadConstant(parent); // array, array, array, index.
        cfw.add(ByteCode.AALOAD); // array, array, parent.
        cfw.addInvoke(
                ByteCode.INVOKESTATIC,
                cfw.getClassName(),
                "init" + index,
                "(Lorg/mozilla/javascript/JSDescriptor;)Lorg/mozilla/javascript/JSDescriptor;"); // array, array, entry.
        cfw.addLoadConstant(index); // array, array, entry, index.
        cfw.add(ByteCode.SWAP); // array, array, index, entry.
        cfw.add(ByteCode.AASTORE);
    }

    private int functionId(JSDescriptor.Builder builder) {
        var code = (OptJSCode.Builder) builder.code;
        return code.index;
    }

    private String mainMethodClassName;
    private CompilerEnvirons compilerEnv;
    private Class<?> targetExtends;
    private Class<?>[] targetImplements;
}
